package server;

import common.*;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * 代理服务器主类
 */
public class ProxyServer {
    private static final Logger logger = LogUtil.getLogger(ProxyServer.class.getName());

    // 配置
    private final ServerConfig config;

    // 客户端连接管理器
    private final ClientManager clientManager;

    // 代理服务管理器
    private final ProxyManager proxyManager;

    // 任务执行器
    private final Executor executor;

    // 定时任务执行器
    private final ScheduledExecutorService scheduler;

    // 控制通道选择器
    private Selector controlSelector;

    // 控制服务器通道
    private ServerSocketChannel controlServerChannel;

    // 运行标志
    private volatile boolean running = false;


    private CountDownLatch latch = new CountDownLatch(1);


    /**
     * 主方法
     */
    public static void main(String[] args) {
        try {
            // 加载配置
            String configPath = args.length > 0 ? args[0] : Thread.currentThread().getContextClassLoader().getResource("config/server.properties").getPath();
            Config configLoader = new Config();
            configLoader.load(configPath);

            // 创建服务器配置
            ServerConfig serverConfig = ServerConfig.fromConfig(configLoader);

            // 创建并启动服务器
            ProxyServer server = new ProxyServer(serverConfig);
            server.start();

            // 添加关闭钩子
            Runtime.getRuntime().addShutdownHook(new Thread(() -> {
                try {
                    server.stop();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }));

        } catch (Exception e) {
            System.err.println("Error starting server: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 构造函数
     */
    public ProxyServer(ServerConfig config) {
        this.config = config;
        this.clientManager = new ClientManager();
        this.proxyManager = new ProxyManager(this);
        this.executor = ConcurrencyUtil.createVirtualThreadExecutor();
        this.scheduler = ConcurrencyUtil.createScheduler(2, "server-scheduler");
    }

    /**
     * 启动代理服务器
     */
    public void start() throws IOException {
        if (running) {
            logger.warning("Server is already running");
            return;
        }

        // 初始化控制通道
        initControlChannel();

        // 启动代理服务
        startProxyServices();

        // 设置定时任务
        setupScheduledTasks();

        // 设置运行标志
        running = true;

        logger.info("Proxy server started on port " + config.getControlPort());

        // 开始接受连接
        executor.execute(this::acceptLoop);

        try {
            latch.await();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 初始化控制通道
     */
    private void initControlChannel() throws IOException {
        controlSelector = Selector.open();
        controlServerChannel = ServerSocketChannel.open();
        controlServerChannel.configureBlocking(false);
        controlServerChannel.socket().bind(new InetSocketAddress(config.getControlPort()));
        controlServerChannel.register(controlSelector, SelectionKey.OP_ACCEPT);
    }

    /**
     * 启动代理服务
     */
    private void startProxyServices() throws IOException {
        ProxyMapping[] mappings = config.getProxyMappings();
        if (mappings != null) {
            for (ProxyMapping mapping : mappings) {
                if (mapping != null) {
                    proxyManager.addProxyService(mapping);
                }
            }
        }
    }

    /**
     * 设置定时任务
     */
    private void setupScheduledTasks() {
        // 心跳检查任务
        scheduler.scheduleAtFixedRate(() -> {
            try {
                clientManager.checkHeartbeats();
            } catch (Exception e) {
                logger.log(Level.SEVERE, "Error in heartbeat check", e);
            }
        }, Protocol.HEARTBEAT_INTERVAL, Protocol.HEARTBEAT_INTERVAL, TimeUnit.MILLISECONDS);

        // 统计信息任务
        scheduler.scheduleAtFixedRate(() -> {
            try {
                logStatistics();
            } catch (Exception e) {
                logger.log(Level.SEVERE, "Error in statistics logging", e);
            }
        }, 60000, 60000, TimeUnit.MILLISECONDS);
    }

    /**
     * 连接接受循环
     */
    private void acceptLoop() {
        try {
            while (running) {
                if (controlSelector.select() > 0) {
                    Iterator<SelectionKey> keys = controlSelector.selectedKeys().iterator();

                    while (keys.hasNext()) {
                        SelectionKey key = keys.next();
                        keys.remove();

                        if (!key.isValid()) {
                            continue;
                        }

                        if (key.isAcceptable()) {
                            acceptConnection();
                        }
                    }
                }
            }
        } catch (Exception e) {
            logger.log(Level.SEVERE, "Error in accept loop", e);
        }
    }

    /**
     * 接受新连接
     */
    private void acceptConnection() {
        try {
            SocketChannel clientChannel = controlServerChannel.accept();
            if (clientChannel != null) {
                clientChannel.configureBlocking(false);

                // 处理新客户端连接
                executor.execute(() -> handleClientConnection(clientChannel));
            }
        } catch (IOException e) {
            logger.log(Level.SEVERE, "Error accepting connection", e);
        }
    }

    /**
     * 处理客户端连接
     */
    private void handleClientConnection(SocketChannel clientChannel) {
        try {
            // 设置TCP配置
            clientChannel.socket().setTcpNoDelay(true);
            clientChannel.socket().setKeepAlive(true);

            // 创建客户端处理器
            ClientHandler handler = new ClientHandler(this, clientChannel);

            // 开始处理客户端通信
            handler.start();

        } catch (Exception e) {
            logger.log(Level.SEVERE, "Error handling client connection", e);
            try {
                clientChannel.close();
            } catch (IOException closeEx) {
                logger.log(Level.WARNING, "Error closing client channel", closeEx);
            }
        }
    }

    /**
     * 停止代理服务器
     */
    public void stop() {
        if (!running) {
            return;
        }

        running = false;

        try {
            // 关闭所有客户端连接
            clientManager.closeAll();

            // 关闭所有代理服务
            proxyManager.closeAll();

            // 关闭控制通道
            if (controlServerChannel != null) {
                controlServerChannel.close();
            }

            if (controlSelector != null) {
                controlSelector.close();
            }

            // 关闭调度器
            ConcurrencyUtil.shutdownExecutor(scheduler);

            logger.info("Proxy server stopped");
        } catch (Exception e) {
            logger.log(Level.SEVERE, "Error stopping server", e);
        } finally {
            latch.countDown();
        }
    }

    /**
     * 记录统计信息
     */
    private void logStatistics() {
        int clientCount = clientManager.getClientCount();
        int connectionCount = proxyManager.getConnectionCount();

        logger.info(String.format("Statistics - Clients: %d, Connections: %d",
                clientCount, connectionCount));
    }

    /**
     * 获取客户端管理器
     */
    public ClientManager getClientManager() {
        return clientManager;
    }

    /**
     * 获取代理管理器
     */
    public ProxyManager getProxyManager() {
        return proxyManager;
    }

    /**
     * 获取配置
     */
    public ServerConfig getConfig() {
        return config;
    }

    /**
     * 获取执行器
     */
    public Executor getExecutor() {
        return executor;
    }


}
